#!/usr/bin/env python
"""
Electronic structure solver.

Type:

$ ./schroedinger.py

for usage and help.
"""
import os
from optparse import OptionParser

import sfepy
from sfepy.base.base import output
from sfepy.base.conf import ProblemConf, get_standard_keywords
from sfepy.base.ioutils import ensure_path
from sfepy.physics.schroedinger_app import SchroedingerApp

def fix_path(filename):
    return os.path.join(sfepy.data_dir, filename)

usage = """%prog [options] [filename_in]

Solver for electronic structure problems.

You need to create a mesh (optionally specify a dimension):

    $ ./schroedinger.py --create-mesh --2d

and then pick a problem to solve, some examples below (the dimension is
determined by the mesh that you created above):

    $ ./schroedinger.py --hydrogen
    $ ./schroedinger.py --well
    $ ./schroedinger.py --boron
    $ ./schroedinger.py --oscillator

and visualize the result:

- using Mayavi

  - 2D:

    $ ./postproc.py mesh.vtk

  - 3D:

    $ ./postproc.py mesh.vtk --3d

- using ParaView

    $ paraview --data=mesh.vtk
"""

help = {
    'conf' :
    'override problem description file items, written as python'
    ' dictionary without surrounding braces',
    'options' : 'override options item of problem description,'
    ' written as python dictionary without surrounding braces',
    'filename' :
    'basename of output file(s) [default: <basename of input file mesh>]',
    'well' :
    'solve infinite potential well (particle in a box) problem',
    'oscillator' :
    'solve spherically symmetric linear harmonic oscillator '
    '(1 electron) problem',
    'hydrogen' :
    'solve the hydrogen atom',
    'boron' :
    'solve the boron atom with 1 electron',
    'create_mesh' :
    'creates a mesh',
    'dim' :
    'Create a 2D mesh, instead of the default 3D',
    'mesh' :
    'override mesh file name. If given in form function(args), mesh is'
    ' generated on the fly.'
    ' Supported functions are cube(edge/2, nodes on edge) and'
    ' sphere(r, nodes density at start, end)',
    'mesh_dir' :
    'directory, where the mesh is generated by --mesh or --create-mesh'
    ' [default: %default]',
}

def main():
    parser = OptionParser(usage=usage, version='%prog ' + sfepy.__version__)
    parser.add_option('-c', '--conf', metavar='"key : value, ..."',
                      action='store', dest='conf', type='string',
                      default=None, help= help['conf'])
    parser.add_option('-O', '--options', metavar='"key : value, ..."',
                      action='store', dest='app_options', type='string',
                      default=None, help=help['options'])
    parser.add_option('-o', '', metavar='filename',
                      action='store', dest='output_filename_trunk',
                      default=None, help=help['filename'])
    parser.add_option('--create-mesh',
                      action='store_true', dest='create_mesh',
                      default=False, help=help['create_mesh'])
    parser.add_option('--2d',
                      action='store_true', dest='dim2',
                      default=False, help=help['dim'])
    parser.add_option('-m', '--mesh', metavar='filename',
                      action='store', dest='mesh',
                      default=None, help=help['mesh'])
    parser.add_option('--mesh-dir', metavar='dirname',
                      action='store', dest='mesh_dir',
                      default='tmp', help=help['mesh_dir'])
    parser.add_option('--oscillator',
                      action='store_true', dest='oscillator',
                      default=False, help=help['oscillator'])
    parser.add_option('--well',
                      action='store_true', dest='well',
                      default=False, help=help['well'])
    parser.add_option('--hydrogen',
                      action='store_true', dest='hydrogen',
                      default=False, help=help['hydrogen'])
    parser.add_option('--boron',
                      action='store_true', dest='boron',
                      default=False, help=help['boron'])

    options, args = parser.parse_args()

    if options.create_mesh and options.mesh:
        output('--create-mesh and --mesh options are mutually exclusive!')
        return

    if len(args) == 1:
        filename_in = args[0];
        auto_mesh_name = False

    elif len(args) == 0:
        auto_mesh_name = True

        mesh_filename = os.path.join(options.mesh_dir, 'mesh.vtk')
        ensure_path(mesh_filename)

        if options.oscillator:
            filename_in = fix_path("examples/quantum/oscillator.py")

        elif options.well:
            filename_in = fix_path("examples/quantum/well.py")

        elif options.hydrogen:
            filename_in = fix_path("examples/quantum/hydrogen.py")

        elif options.boron:
            filename_in = fix_path("examples/quantum/boron.py")

        elif options.create_mesh:
            output('generating mesh...')
            try:
                os.makedirs("tmp")
            except OSError, e:
                if e.errno != 17: # [Errno 17] File exists
                    raise
            if options.dim2:
                output("dimension: 2")
                gp = fix_path('meshes/quantum/square.geo')
                os.system("cp %s tmp/mesh.geo" % gp)
                os.system("gmsh -2 tmp/mesh.geo -format mesh")
                mtv = fix_path('script/convert_mesh.py')
                os.system("%s tmp/mesh.mesh %s" % (mtv, mesh_filename))
            else:
                output("dimension: 3")
                import sfepy.mesh as geom
                from sfepy.fem.mesh import Mesh
                try:
                    from site_cfg import tetgen_path
                except ImportError:
                    tetgen_path = '/usr/bin/tetgen'
                gp = fix_path('meshes/quantum/box.geo')
                os.system("gmsh -0 %s -o tmp/x.geo" % gp)
                g = geom.read_gmsh("tmp/x.geo")
                g.printinfo()
                geom.write_tetgen(g, "tmp/t.poly")
                geom.runtetgen("tmp/t.poly", a=0.03, Q=1.0,
                               quadratic=False, tetgenpath=tetgen_path)
                m = Mesh.from_file("tmp/t.1.node")
                m.write(mesh_filename, io="auto")
            output("...mesh written to %s" % mesh_filename)
            return

        else:
            parser.print_help()
            return

    else:
        parser.print_help()
        return

    required, other = get_standard_keywords()
    conf = ProblemConf.from_file_and_options(filename_in, options,
                                             required, other)

    if options.mesh:
        from sfepy.fem.mesh_generators import gen_mesh_from_string

        conf.filename_mesh = gen_mesh_from_string(options.mesh,
                                                  options.mesh_dir)

    elif auto_mesh_name and not sfepy.in_source_tree:
        conf.filename_mesh = mesh_filename
        conf.options.absolute_mesh_path = True

    app = SchroedingerApp(conf, options, 'schroedinger:')
    opts = conf.options
    if hasattr(opts, 'parametric_hook'): # Parametric study.
        parametric_hook = conf.get_function(opts.parametric_hook)
        app.parametrize(parametric_hook)
    app()

if __name__ == '__main__':
    main()
